package com.learn.window.aggregation;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.ReducingState;
import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.Iterator;

/**
 * @create: 2023-04-24 00:09
 * @author: Mr.Du
 * --------------
 * @notes:
 * 1000, spark, 2
 * 5000,spark,2
 * 6000,spark,3
 * 10000,spark,5
 * TODO 可能有问题
 **/
public class WindowStateAndGlobalStateFunctionDemo {
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.enableCheckpointing(10000);

        DataStreamSource<String> lines = env.socketTextStream("node1", 9999);

        WindowedStream<Tuple2<String, Integer>, String, TimeWindow> window = lines
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<String>forBoundedOutOfOrderness(Duration.ZERO)
                                .withTimestampAssigner((line, timestamp) -> Long.parseLong(line.split(",")[0]))
                )
                .map(new MapFunction<String, Tuple2<String, Integer>>() {
                    @Override
                    public Tuple2<String, Integer> map(String value) throws Exception {
                        String[] fields = value.split(",");
                        return Tuple2.of(fields[1], Integer.parseInt(fields[2]));
                    }
                }).keyBy(t -> t.f0)
                .window(TumblingEventTimeWindows.of(Time.seconds(5)));

        //需求：可以在窗口内进行增量聚合，并且还可以与历史数据进行聚合
        window
                .aggregate(new AggregateFunction<Tuple2<String, Integer>, Integer, Integer>() {
                    //创建一个初始值
                    @Override
                    public Integer createAccumulator() {
                        return 0;
                    }

                    @Override
                    public Integer add(Tuple2<String, Integer> value, Integer accumulator) {
                        return value.f1 + accumulator;
                    }

                    @Override
                    public Integer getResult(Integer accumulator) {
                        return accumulator;
                    }

                    //如果不是sessionWindow可以不实现
                    @Override
                    public Integer merge(Integer integer, Integer acc1) {
                        return null;
                    }
                }, new ProcessWindowFunction<Integer, String, String, TimeWindow>() {
                    private transient ReducingStateDescriptor<Integer> windowStateDescriptor;
                    private transient ReducingStateDescriptor<Integer> globalStateDescriptor;
                    @Override
                    public void open(Configuration parameters) throws Exception {
                        windowStateDescriptor = new ReducingStateDescriptor<Integer>("window", new ReduceFunction<Integer>() {
                            @Override
                            public Integer reduce(Integer v1, Integer v2) throws Exception {
                                return v1 + v2;
                            }
                        }, Integer.class);
                        globalStateDescriptor = new ReducingStateDescriptor<Integer>("global", new ReduceFunction<Integer>() {
                            @Override
                            public Integer reduce(Integer v1, Integer v2) throws Exception {
                                return v1 + v2;
                            }
                        }, Integer.class);
                    }

                    @Override
                    public void process(String key, ProcessWindowFunction<Integer, String, String, TimeWindow>.Context context, Iterable<Integer> iterable, Collector<String> out) throws Exception {
                        Integer sum = 0;
                        for (Integer integer : iterable) {
                            sum += integer;
                        }
                        ReducingState<Integer> windowState = context.windowState()
                                .getReducingState(windowStateDescriptor);
                        ReducingState<Integer> globalState = context.windowState()
                                .getReducingState(globalStateDescriptor);

                        iterable.forEach(t -> {
                            try {
                                windowState.add(t);
                                globalState.add(t);
                            } catch (Exception e) {
                                e.printStackTrace();
                            }
                        });
                        out.collect(key+",window:" + windowState.get()+", global:" + globalState.get());
                    }
                }).print();

        env.execute();

    }
}
